package cs429_project.clustering;

import java.io.PrintStream;

import java.math.BigDecimal;
import java.math.RoundingMode;

import java.util.*;

/**
 * Represents a cluster of documents for k-means clustering.  Document vectors
 * can be added and removed, and the mean of the cluster can be computed.  For
 * this class to work properly, all document vectors must contain mappings for
 * all of the same terms; any differences may trigger errors.
 * @see cs429_project.clustering.KMeansClustering
 * @author Peter Ferrans
 */
public class DocCluster {

    /** Set of all documents in the cluster. */
    private TreeSet<DocVec> docSet;
    /** The mean for this cluster.  Maps each term in the cluster to its
     * average frequency. */
    private TreeMap<String, BigDecimal> mean;

    /**
     * Creates a document cluster with no documents.
     */
    public DocCluster() {
        this.docSet = new TreeSet<DocVec>();
        this.mean = computeMean();
    }

    /**
     * Creates a cluster containing all of the documents passed in through the
     * parameter.
     * @param docVecs
     *    Vector representations of documents for this cluster.
     */
    public DocCluster(Set<DocVec> docVecs) {
        this.docSet = new TreeSet<DocVec>(docVecs);
        this.mean = computeMean();
    }

    /**
     * Retrieves the document vectors in this cluster.
     * @return
     *    A set of the document vectors in this cluster.
     */
    public TreeSet<DocVec> getDocVecSet() {
        return this.docSet;
    }

    /**
     * Retrieves the document vector in this cluster with the specified ID
     * number.
     * @param docId
     *    ID number of the document vector to retrieve.
     * @return
     *    The document vector with the specified ID number, or {@code null} if
     *    this cluster does not contain a document vector with that ID.
     */
    public DocVec getDocVec(int docId) {
        DocVec dv = null;
        Iterator<DocVec> it = this.docSet.iterator();
        while ((it.hasNext()) && (dv == null)) {
            DocVec candDv = it.next();
            if (docId == candDv.docId) {
                dv = candDv;
            }
        }

        return dv;
    }

    /**
     * Adds a document vector to the cluster.
     * @param docVec
     *    Document vector to be added.
     * @return
     *    Whether or not the addition was successful.
     */
    public boolean addDocVec(DocVec docVec) {
        return this.docSet.add(docVec);
    }

    /**
     * Removes a document vector from the cluster.
     * @param docVec
     *    Document vector to be removed.
     * @return
     *    Whether or not the removal was successful.
     */
    public boolean removeDocVec(DocVec docVec) {
        return this.docSet.remove(docVec);
    }

    /**
     * Removes from the cluster the document vector with the specified ID.
     * @param docId
     *    ID number for the document vector to be removed.
     * @return
     *    The document vector that was removed.  If no document vector existed
     *    in the collection with the specified ID, {@code null} is returned.
     */
    public DocVec removeDocVec(int docId) {
        DocVec dv = getDocVec(docId);
        if (dv != null) {
            this.docSet.remove(dv);
        }

        return dv;
    }

    /**
     * Removes all documents from the cluster.
     */
    public void clear() {
        this.docSet.clear();
    }

    /**
     * Retrieves the current mean of the cluster: a map of the average
     * frequencies of all terms that occur in documents in the cluster.
     * @return
     *    Current mean of the cluster.
     */
    public TreeMap<String, BigDecimal> getCurrentMean() {
        return this.mean;
    }

    /**
     * Recalculates the current mean of the cluster.  The mean is not
     * recalculated immediately whenever a document vector is added or removed,
     * because for k-means clustering the mean needs to stay the same while
     * documents are being moved around from one cluster to another.
     */
    public void resetMean() {
        if(this.docSet.isEmpty()) {
            this.mean = null;
        } else {
            this.mean = computeMean();
        }
    }

    /**
     * Computes the maximum likelihood estimate value for a term in this
     * cluster.
     * @param term
     *    Term for which we want the maximum likelihood estimate.
     * @return
     *    Maximum likelihood estimate for the term.
     */
    public BigDecimal maxLikeEst(String term) {
        BigDecimal mle = BigDecimal.ZERO;
        Iterator<DocVec> docIt = this.docSet.iterator();
        while(docIt.hasNext()) {
            DocVec dv = docIt.next();
            BigDecimal thisFreq = dv.get(term);
            if(thisFreq != null) {
                mle = mle.add(thisFreq);
            }
        }
        BigDecimal size = numberOfTerms();
        mle = mle.divide(size, DocVec.DEFAULT_PRECISION, RoundingMode.UP);

        return mle;
    }

    /**
     * Computes the total number of words in this cluster.
     * @return
     *    Total number of words in the cluster.
     */
    public BigDecimal numTotalWords() {
        BigDecimal numTotalWords = BigDecimal.ZERO;
        Iterator<DocVec> dvIt = this.docSet.iterator();
        while(dvIt.hasNext()) {
            DocVec docVec = dvIt.next();
            numTotalWords = numTotalWords.add(docVec.numTotalWords());
        }

        return numTotalWords;
    }

    /**
     * Computes the frequency of a word in all documents in this cluster.
     * @param word
     *    Word whose frequency is to be computed.
     * @return
     *    The frequency of the word in this cluster.
     */
    public BigDecimal freqInCluster(String word) {
        BigDecimal freq = BigDecimal.ZERO;
        Iterator<DocVec> dvIt = this.docSet.iterator();
        while(dvIt.hasNext()) {
            DocVec docVec = dvIt.next();
            BigDecimal thisFreq = docVec.get(word);
            if(thisFreq != null) {
                freq = freq.add(thisFreq);
            }
        }

        return freq;
    }

    /**
     * Computes for each word the total number of times the word occurs in
     * documents in the cluster.
     * @return
     *    A mapping of each word to the number of times it appears.
     */
    public TreeMap<String, BigDecimal> computeSums() {
        TreeMap<String, BigDecimal> sumMap = new TreeMap<String, BigDecimal>();
        TreeSet<String> termSet = getAllTerms();
        Iterator<String> termIt = termSet.iterator();
        while(termIt.hasNext()) {
            String term = termIt.next();
            BigDecimal sum = termSum(term);
            sumMap.put(term, sum);
        }

        return sumMap;
    }

    /**
     * Expands the terms in each document vector.  This means that each
     * document vector will contain an entry for every word that occurs in any
     * document in this cluster.  This allows us to compare document vectors to
     * one another.
     */
    public final void expandDocTerms() {
        TreeSet<String> terms = getAllTerms();
        expandDocTerms(terms);
    }

    /**
     * Prints the data for each document vector to the specified print stream.
     * @param stream
     *    Stream to which document information will be printed.
     */
    public void printDocs(PrintStream stream) {
        Iterator<DocVec> docSetIt = this.docSet.iterator();
        while(docSetIt.hasNext()) {
            DocVec dv = docSetIt.next();
            dv.print(stream);
            stream.println();
        }
    }

    /**
     * Computes the total number of unique terms found in this cluster.  This
     * differs from {@link #numTotalWords()}, which merely counts the total
     * number of words appearing.
     * @return
     *    The number of unique terms found in this cluster.
     */
    public BigDecimal numberOfTerms() {
        TreeSet<String> terms = getAllTerms();
        BigDecimal numTerms = BigDecimal.valueOf((double)(terms.size()));

        return numTerms;
    }

    /**
     * Creates a set of all the terms that occur anywhere in this cluster.
     * @return
     *    Set of all terms in this cluster.
     */
    protected TreeSet<String> getAllTerms() {
        TreeSet<String> termSet = new TreeSet<String>();
        Iterator<DocVec> it = this.docSet.iterator();
        while (it.hasNext()) {
            DocVec docVec = it.next();
            termSet.addAll(docVec.getTerms());
        }

        return termSet;
    }

    /**
     * Given a set of terms, expands all document vectors in this cluster so
     * that they have entries for each of the terms in the set.
     * @param terms
     *    Set of terms with which to expand document vectors in this cluster.
     */
    protected void expandDocTerms(Set<String> terms) {
        TreeSet<DocVec> expDocSet = new TreeSet<DocVec>();
        Iterator<DocVec> it = this.docSet.iterator();
        while (it.hasNext()) {
            DocVec docVec = it.next();
            docVec.expand(terms);
            expDocSet.add(docVec);
        }

        this.docSet = expDocSet;
    }

    /**
     * Computes the mean of the documents in this cluster.  Used as a helper
     * function for {@link #resetMean()} and the constructors.
     * @return
     *    Mean map for the terms in this cluster.
     */
    private TreeMap<String, BigDecimal> computeMean() {
        TreeMap<String, BigDecimal> meanMap = computeSums();
        BigDecimal size = BigDecimal.valueOf(this.docSet.size());
        Iterator<String> termIt = meanMap.keySet().iterator();

        while(termIt.hasNext()) {
            String term = termIt.next();
            BigDecimal termSum = meanMap.get(term);
            BigDecimal termMean = termSum.divide(size,
                    DocVec.DEFAULT_PRECISION, RoundingMode.UP);
            meanMap.put(term, termMean);
        }

        return meanMap;
    }

    /**
     * Computes the total number of times a term appears in documents in this
     * cluster.
     * @param term
     *    Term whose occurrences we are counting.
     * @return
     *    Number of occurrences of the term in this collection.
     */
    private BigDecimal termSum(String term) {
        Iterator<DocVec> docIt = this.docSet.iterator();
        BigDecimal sum = BigDecimal.ZERO;
        while(docIt.hasNext()) {
            DocVec docVec = docIt.next();
            BigDecimal thisFreq = docVec.get(term);
            sum = sum.add(thisFreq);
        }

        return sum;
    }
}
